import { Cell, CodeCell } from '@jupyterlab/cells';
import { ToolbarButton} from '@jupyterlab/apputils';
import { Widget } from '@lumino/widgets';
import { INotebookTracker,NotebookPanel, Notebook, NotebookActions} from '@jupyterlab/notebook';
import NotebookUtilities from "./utils"
import { Constants } from './constants';


class JupyterlabCodeOptimizer {
    protected working: boolean;
    protected panel :NotebookPanel;
    private tmp_path: string;
    public log_path:string;
    public tmp_log_path:string;
    public rand:number
    public markdown: Cell | undefined    
    public cells:CodeCell[];
    constructor(panel: NotebookPanel) {
      this.working = false;
      this.panel = panel
      this.tmp_path = Constants.WORK_PATH+"tmp.py"
      this.rand = NotebookUtilities.GetRandomNum(0,200)
      this.log_path = Constants.WORK_PATH + "NeuralCoder"+this.rand+".log";
      this.tmp_log_path = Constants.WORK_PATH+ "NeuralCoder_tmp"+".log";
      this.cells = [];
    }
 

     async optimizeCode(
      code: string[],
      formatter: string,
      name:string,
      next:string,
      options: string | undefined,
      notebook: boolean,
      panel: NotebookPanel,
      cell:CodeCell,
      run?:ToolbarButton|undefined,
    ){
      let codes: string[] = []
      code.forEach(function (value) {
        value = value.replace(/('\\n')/g, '^^^')
        value = value.replace(/\\n"/g, '###')
        value = value.replace(/\\n'/g, '###')
        value = value.replace(/"\\n/g, '@@')
        value = value.replace(/'\\n/g, '@@')
        value = value.replace(/\n/g, '\\n')
        value = value.replace(/"/g, '+++')
        value = value.replace(/,/g, '$')        
        codes.push(value)    
      })
      
      let gen_code = `code = "${codes}"\ncodes = code.split(',')\nwith open( '${this.tmp_path}', 'w+' ) as f:\n    for i in range(0,len(codes)):\n        f.write('# this is the beginning of a single code snippet\\n')\n        code_list = codes[i].replace('$',',').replace('+++','\"').split('\\n')\n        for line in code_list:\n            if('split(^^^)' in line):\n                    line=line.replace('split(^^^)', 'split(\\'\\\\n\\')')\n            if('###' in line):\n                    line=line.replace('###', '\\\\n\"')\n            if('@@' in line):\n                    line=line.replace('@@', '\"\\\\n')\n            f.write(line+'\\n')`
      const expr = {code_list: `code_list`};
      NotebookUtilities.sendKernelRequestFromNotebook(panel, gen_code, expr,false);

      if(options === 'normal'){
        let runcode = `from neural_coder import enable\nenable(code="${this.tmp_path}",features=["${formatter}"], overwrite=True)`;
        let expr = {sum: ` `};
        NotebookUtilities.sendKernelRequestFromNotebook(panel, runcode, expr,false);      
        let run_code1 = `with open("${this.tmp_path}", 'r') as f:\n    optimized_code = f.read()\n`
        let expr1 = {optimizedCode : "optimized_code"}
        let result2 = NotebookUtilities.sendKernelRequestFromNotebook(panel,run_code1,expr1,false)
        result2.then(value=>{
          let optimizedTexts = Object.values(value.optimizedCode.data)[0] as string 
          let optimizeCodes = optimizedTexts.split('# this is the beginning of a single code snippet\\n').slice(1,)
          optimizeCodes[optimizeCodes.length-1] = optimizeCodes[optimizeCodes.length-1].slice(0,-3)
          for (let i = 0; i < optimizeCodes.length; ++i) {
            const cell = this.cells[i];
            const currentTexts = this.cells.map(cell => cell.model.value.text);
            const currentText = currentTexts[i];
            let optimizedtext = optimizeCodes[i];
            optimizedtext = optimizedtext.replace(/\\'\\\\n\\'/g, "^^^")
            optimizedtext = optimizedtext.replace(/\\\\n"/g, "+++")
            optimizedtext = optimizedtext.replace(/\\\\n'/g, "+++")
            optimizedtext = optimizedtext.replace(/"\\\\n/g, "@@@")
            optimizedtext = optimizedtext.replace(/'\\\\n/g, "@@@")
            optimizedtext = optimizedtext.replace(/\\n/g,'\n')
            optimizedtext = optimizedtext.replace(/\\'/g,"'")
            optimizedtext = optimizedtext.replace(/\^\^\^/g, "'\\n'")
            optimizedtext = optimizedtext.replace(/\+\+\+/g, "\\n\"")
            optimizedtext = optimizedtext.replace(/\@\@\@/g, "\"\\n")
            if (cell.model.value.text === currentText) {
                cell.model.value.text = optimizedtext;
              }
            const run_svg = document.createElement("svg")
            run_svg.innerHTML = Constants.ICON_RUN
            run?.node.firstChild?.firstChild?.firstChild?.firstChild?.replaceWith(run_svg)
            } 
        })
       
      }else{
        if(formatter === ''){
            if (this.markdown){
              this.markdown.model.value.text += "[NeuralCoder INFO] Enabling and Benchmarking for The Original Model ......  \n"
            }
            let runcode1 = `with open("${this.log_path}", 'a' ) as f:\n       f.write("[NeuralCoder INFO] Enabling and Benchmarking for The Original Model ......\\n")`
            let expr1 = {path: ""};
            NotebookUtilities.sendKernelRequestFromNotebook(panel, runcode1, expr1,false);

            let runcode = `from neural_coder import enable\nperfomance, mode, path = enable(code="${this.tmp_path}",features=[], run_bench=True, args="${options}")\nwith open(path + '/bench.log', 'r') as f:\n    logs = f.readlines()\nlog_line = logs[4]\nlog = log_line.split("[")[1].split("]")[0]`
            let expr = {path: "path",log:"log"};
            let result = NotebookUtilities.sendKernelRequestFromNotebook(panel, runcode, expr,false);
            let fps : String ;
            result.then(  value => {
              fps =  Object.values(value.log.data)[0] as string;
              if (this.markdown){
                this.markdown.model.value.text += `[NeuralCoder INFO] Benchmark Result (Performance) of The Original Model is ${fps} (FPS)  \n`
              }
              let text = `[NeuralCoder INFO] Benchmark Result (Performance) of The Original Model is ${fps} (FPS)\\n`
              let runcode = `with open("${this.log_path}", 'a' ) as f:\n   f.write("${text}")`
              let expr = {path: ""};
              NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode, expr,false);          
              if (this.markdown){
                this.markdown.model.value.text += `[NeuralCoder INFO] Enabling and Benchmarking for ${next} ......  \n`
              }
              let runcode1 = `with open("${this.log_path}", 'a' ) as f:\n       f.write("[NeuralCoder INFO] Enabling and Benchmarking for ${next} ......\\n")`
              let expr1 = {path: ""};
              NotebookUtilities.sendKernelRequestFromNotebook(panel, runcode1, expr1,false);
              let runcode2 = `with open("${this.tmp_log_path}", 'a' ) as f:\n       f.write("${text}")`
              let expr2 = {path: ""};
              NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode2, expr2,false);
            })
          
        }else{

          let runcode = `from neural_coder import enable\nperfomance, mode, path = enable(code="${this.tmp_path}", features=["${formatter}"], run_bench=True, args="${options}")\nwith open(path + '/bench.log', 'r') as f:\n    logs = f.readlines()\nlog_line = logs[4]\nlog = log_line.split("[")[1].split("]")[0]`
          let expr = {path: "path",log:"log"};
          let result = NotebookUtilities.sendKernelRequestFromNotebook(panel, runcode, expr,false);
          let fps:string;
          result.then( value => {
            fps = Object.values(value.log.data)[0] as string;
            if(this.markdown){
              this.markdown.model.value.text += `[NeuralCoder INFO] Benchmark Result (Performance) of ${name} is ${fps} (FPS)  \n`
            }
            let text = `[NeuralCoder INFO] Benchmark Result (Performance) of ${name} is ${fps} (FPS)\\n`
            let runcode = `with open("${this.log_path}", 'a' ) as f:\n       f.write("${text}")`
            let expr = {path: ""};
            NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode, expr,false);
            if (next !== ''){
              if(this.markdown){
                this.markdown.model.value.text += `[NeuralCoder INFO] Enabling and Benchmarking for ${next} ......  \n`
              }
              let runcode2 = `with open("${this.log_path}", 'a' ) as f:\n       f.write("[NeuralCoder INFO] Enabling and Benchmarking for ${next} ......\\n")`
              let expr2 = {path: ""};
              NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode2, expr2,false);
            }
            let runcode2 = `with open("${this.tmp_log_path}", 'a' ) as f:\n       f.write("${text}")`
            let expr2 = {path: ""};
            NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode2, expr2,false);
            
            if (formatter === 'pytorch_aliblade'){
              let read_log = `import re\nwith open("${this.tmp_log_path}", 'r') as f:\n    logs = f.readlines()\n    fps_list=[]\n    for log_line in logs[-2:]:\n        pat = r\'\\d+\\.?\\d+'\n        fps = re.search(pat,log_line).group()\n        fps_list.append(float(fps))\nmaxi = max(fps_list)\nindex = fps_list.index(maxi)\nboost = round(maxi/fps_list[0],1)\nfeatures=['','pytorch_aliblade',]\nfeature_name=['','Alibaba Blade-DISC']\nbest_feature = features[index]\nbest_name = feature_name[index]\n`
              let read_expr = {boost:"boost",best_feature:"best_feature",best_name:"best_name"}
              let read_result = NotebookUtilities.sendKernelRequestFromNotebook(this.panel, read_log, read_expr,false);
              read_result.then(value =>{
                let boost = Object.values(value.boost.data)[0] as string;
                let best_name = Object.values(value.best_name.data)[0] as string;
                if(this.markdown){
                  this.markdown.model.value.text += `[NeuralCoder INFO] The Best Intel Optimization: ${best_name}  \n`
                  this.markdown.model.value.text += `[NeuralCoder INFO] You can get up to ${boost}X performance boost.  \n`
                }   
                  let command = "lscpu | grep 'Model name'"
                  let get_hardware = `import subprocess\nsubp = subprocess.Popen("${command}",shell=True,stdout=subprocess.PIPE,stderr=subprocess.PIPE,encoding="utf-8")\nsubp.wait(2)\nhardware = subp.communicate()[0].replace("Model name:","").strip()`
                  let expr_hardware = {hardware: "hardware"}
                  let hard_res = NotebookUtilities.sendKernelRequestFromNotebook(this.panel, get_hardware, expr_hardware,false);
                  hard_res.then(value =>{
                    let hard = Object.values(value.hardware.data)[0] as string;
                    if(this.markdown){
                      this.markdown.model.value.text += `[NeuralCoder INFO] HardWare: ${hard}  \n`
                      this.markdown.model.value.text +=  `[NeuralCoder INFO] The log was saved to neural_coder_workspace\\NeuralCoder${this.rand}.log  \n`
                    }
                   })
                   const run_svg = document.createElement("svg")
                  run_svg.innerHTML = Constants.ICON_RUN
                  run?.node.firstChild?.firstChild?.firstChild?.firstChild?.replaceWith(run_svg)
                  })
              } 
          })
        }
      }  
    }
  }
  

  export class JupyterlabNotebookCodeOptimizer extends JupyterlabCodeOptimizer {
    protected notebookname: string;
    protected notebookTracker: INotebookTracker;

    constructor(
      notebookTracker: INotebookTracker,
      panel: NotebookPanel
    ) {
      super(panel);
      this.notebookTracker = notebookTracker;
      this.notebookname = '';
    }
  
    public async optimizeAction(config: any, formatter?: string) {
      return this.optimizeCells(true, config, formatter);
    }
  
  
    public async optimizeAllCodeCells(
      config?: string,
      formatter?: string,
      notebook?: Notebook,
      run?: ToolbarButton
    ):Promise<string> {
      return this.optimizeCells(false, config, formatter, notebook, run);
    }

    private getCodeCells(ifMarkdown = true, notebook?: Notebook): CodeCell[] {
      if (!this.notebookTracker.currentWidget) {
        return [];
      }
      const codeCells: CodeCell[] = [];
      notebook = notebook || this.notebookTracker.currentWidget.content;      
      this.notebookname = notebook.title.label
      let count = 0;
      notebook.widgets.forEach((cell: Cell) => {
        if (cell.model.type === 'code') {
          count += 1;
          codeCells.push(cell as CodeCell);
        }
      });
      
      if(ifMarkdown){
        NotebookActions.insertBelow(notebook);
        this.notebookTracker.currentWidget.content.activeCellIndex = count + 1;
        NotebookActions.changeCellType(notebook, 'markdown')
        const activeCell = notebook.activeCell;
        if (activeCell){
          this.markdown = activeCell
      }  
      }
      
      
      this.cells = codeCells
      return codeCells;
    }
  
    private async optimizeCells(
      selectedOnly: boolean,
      config?: string,
      formatter?: string,
      notebook?: Notebook,
      run?: ToolbarButton|undefined
    ):Promise<string>{
      if (this.working) {
        return  new Promise<string>((resolve, reject) => {
          resolve("false!")
        });
      }
        this.working = true;
        const optimize_type = formatter !== undefined ? formatter : 'pytorch_aliblade';
        if(optimize_type === 'auto'){
          selectedOnly = true
        }else{selectedOnly = false}

        const selectedCells = this.getCodeCells(selectedOnly, notebook);
        let cell = selectedCells[selectedCells.length-1]
        if (selectedCells.length === 0) {
          this.working = false;
          return  new Promise<string>((resolve, reject) => {
            resolve("false!")
          });
        }
        const currentTexts = selectedCells.map(cell => cell.model.value.text);
       if (optimize_type === 'auto'){
        
        if (this.markdown){
          this.markdown.model.value.text += `[NeuralCoder INFO] Code: User code from Jupyter Lab notebook "${this.notebookname}"  \n`
          this.markdown.model.value.text += `[NeuralCoder INFO] Benchmark Mode: Throughput  \n`
        }
        let runcode2 = `with open('${this.log_path}', 'a' ) as f:\n       f.write("[NeuralCoder INFO] Code: User code from Jupyter Lab notebook '${this.notebookname}'\\n")`
        let expr2 = {path: ""};
        NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode2, expr2,false);
        let runcode3 = `with open('${this.log_path}', 'a' ) as f:\n       f.write("[NeuralCoder INFO] Benchmark Mode: Throughput\\n")`
        let expr3 = {path: ""};
        NotebookUtilities.sendKernelRequestFromNotebook(this.panel, runcode3, expr3,false);
        await this.optimizeCode(
              currentTexts,
              '',
              'The Original Model',
              'Alibaba-Blade',
              config,
              true,
              this.panel,
              cell,
              run              
              );
        await this.optimizeCode(
            currentTexts,
            'pytorch_aliblade',
            'Alibaba-Blade',
            '',
            config,
            true,
            this.panel,
            cell,
            run,
            );     

       }else{
        await this.optimizeCode(
          currentTexts,
          optimize_type,
          "",
          "",
          "normal",
          true,
          this.panel,
          cell,
          run,
          );
        }
        this.working = false;
        return  new Promise<string>((resolve, reject) => {
          resolve("success!")
        });
      }
    applicable(formatter: string, currentWidget: Widget) {
      const currentNotebookWidget = this.notebookTracker.currentWidget;
      return currentNotebookWidget && currentWidget === currentNotebookWidget;
    }
  
  }
 
  